# -*- coding: utf-8 -*-
# @Author: lidongdong
# @time  : 19-1-23 下午9:46
# @file  : main.py

import json
import argparse
from utils import *

parser = argparse.ArgumentParser()
parser.add_argument("--mode", default="train", help="choose train, val, test")
parser.add_argument("--json_path", type=str, help="original json annotations file")
parser.add_argument("--image_path", type=str, help="the images path")
parser.add_argument("--new_json_path", type=str, help="the new annotations file")
parser.add_argument("--h5py_path", type=str, help="h5py file storage path")


def main():
    # parse argments
    args = parser.parse_args()
    mode = args.mode
    json_path = args.json_path
    image_path = args.image_path
    new_json_path = args.new_json_path
    h5py_path = args.h5py_path

    print "-" * 20 + "OPTIONS" + "-" * 20
    print vars(args)
    for (k, v) in sorted(vars(args).items()):
        print "{}: {}".format(k, v)

    # prepare new_json_file
    json_content, caption_body = prepare_new_json(json_path, mode, new_json_path)

    # image to h5 file
    h5py_file = json_content["h5py_file"]   # mode + ".h5"
    h5py_file = os.path.join(h5py_path, h5py_file)
    filenames = [caption["filename"] for caption in caption_body]
    image2h5(filenames, h5py_file, root=image_path)
    print "{} DONE".format(mode)


if __name__ == '__main__':
    main()
